import sys
import os

# 添加 project/ 目录到 sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import json
import os
import os.path as osp
import argparse
import cv2
import numpy as np
from PIL import Image
from tqdm import tqdm
import multiprocessing
from tictoc import TicToc
import shutil
from collections import Counter


def check_equal_file_counts_in_subfolders(folder_path):
    # 获取指定文件夹下的所有子文件夹
    subfolders = [os.path.join(folder_path, d) for d in os.listdir(folder_path) if
                  os.path.isdir(os.path.join(folder_path, d))]

    # 初始化一个字典来存储每个子文件夹的文件数量
    file_counts = {}

    # 遍历每个子文件夹并计算文件数量
    for subfolder in subfolders:
        # 获取子文件夹中的文件数量（不包括子文件夹）
        count = sum(1 for _ in os.scandir(subfolder) if _.is_file())
        # 将文件数量与子文件夹路径关联起来
        file_counts[subfolder] = count

    # 使用Counter来计算每种文件数量的出现次数
    count_freq = Counter(file_counts.values())

    # 检查是否只有一种文件数量出现，且其出现次数等于子文件夹的数量
    return len(count_freq) == 1 and count_freq[file_counts[subfolders[0]]] == len(subfolders)


def main_worker(dataset_path, args):
    dataset_ok=True
    camera_dir = os.listdir(os.path.join(dataset_path, 'camera'))
    for subdir in camera_dir:
        if subdir[-6:] == 'normal':
            normal_link = os.path.join(dataset_path, 'camera', subdir)
            if os.path.islink(normal_link):
                dataset_ok=False
                break

    assert dataset_ok
    assert check_equal_file_counts_in_subfolders(os.path.join(dataset_path, 'camera'))

    if os.path.exists(os.path.join(dataset_path, 'camera_epe')):
        shutil.rmtree(os.path.join(dataset_path, 'camera_epe'))
    if os.path.exists(os.path.join(dataset_path, 'camera_extra')):
        shutil.rmtree(os.path.join(dataset_path, 'camera_extra'))

    print(f'done')


def main(args):
    # 统计耗时
    cost = TicToc("carla2real_after")
    assert os.path.exists(args.data_path)
    frames = os.listdir(args.data_path)
    frames.sort(key=lambda x: x)
    files = []
    for dir in frames:
        if dir[:2] == '__':  # '__'
            dir = os.path.join(args.data_path, dir)
            if os.path.isdir(dir):
                files.append(dir)

    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/media/adt/T71/ZWH/docker/files/data/tem/datasets",
                        help='your data root for kitti')
    args = parser.parse_args()

    main(args)
